# -*- coding: UTF-8 -*-
"""
@项目名称：get_file.py
@作   者：陆地起飞全靠浪
@创建日期：2025-09-15-09:21
"""
import json
import os
import os.path
from glob import glob
import gradio as gr
import shutil
from datetime import datetime
import pandas as pd
from rabbit_mq.rmq_config import logger, cfg
from ChemTools.utiltools import ligand_format_conversion, get_verify_smiles
from ChemTools.get_chain_class import get_chain_class
from ChemTools.pdb_add_h import add_hydrogens
from ChemTools.get_pocket import get_pocket
from rabbit_mq.rmq_utiltools import rmq_sent


def test_create_task_id():
    # 获取当前时间
    now = datetime.now()
    # 格式化输出
    formatted_time = now.strftime("%Y_%m_%d-%H_%M_%S")
    task_id = f'generate_{formatted_time}'
    os.makedirs(os.path.join(cfg['generate_result_path'], task_id), exist_ok=True)
    return task_id


def copy_input_file2out(task_id, protein_file=None, smiles_file=None, cqu_pocket_10A_file=None, cqu_ado_sdf_file=None, cqu_ado_smi_file=None):
    '''
    shutil.copyfile 目标地址必须是完整的文件路径，而不能是目标目录路径
    shutil.copy 目标即可以是文件路径也可以是目录路径
    shutil.copy2 的操作和shutil.copy一样，不同的是会拷贝元数据，即拷贝的文件的创建时间、修改时间等和源文件一样，而shutil.copy拷贝的文件的创建时间是新的，和源数据不同。
    '''
    dst_protein_file = None
    if protein_file:  # 处理上传文件
        dst_protein_file = os.path.join(cfg['generate_result_path'], task_id, os.path.basename(protein_file).lower())
        os.makedirs(os.path.dirname(dst_protein_file), exist_ok=True)
        shutil.copyfile(protein_file, dst_protein_file)
    dst_smiles_file = None
    if smiles_file:  # 处理上传文件
        dst_smiles_file = os.path.join(cfg['generate_result_path'], task_id, os.path.basename(smiles_file).lower())
        os.makedirs(os.path.dirname(dst_smiles_file), exist_ok=True)
        shutil.copyfile(smiles_file, dst_smiles_file)
    # 重大ado
    dst_cqu_pocket_10A_file = None
    if cqu_pocket_10A_file:
        dst_cqu_pocket_10A_file = os.path.join(cfg['generate_result_path'], task_id, 'cqu_ado', os.path.basename(cqu_pocket_10A_file).lower())
        os.makedirs(os.path.dirname(dst_cqu_pocket_10A_file), exist_ok=True)
        shutil.copyfile(cqu_pocket_10A_file, dst_cqu_pocket_10A_file)
    dst_cqu_ado_sdf_file = None
    if cqu_ado_sdf_file:
        dst_cqu_ado_sdf_file = os.path.join(cfg['generate_result_path'], task_id, 'cqu_ado', os.path.basename(cqu_ado_sdf_file).lower())
        os.makedirs(os.path.dirname(dst_cqu_ado_sdf_file), exist_ok=True)
        shutil.copyfile(cqu_ado_sdf_file, dst_cqu_ado_sdf_file)
    dst_cqu_ado_smi_file = None
    if cqu_ado_smi_file:
        dst_cqu_ado_smi_file = os.path.join(cfg['generate_result_path'], task_id, 'cqu_ado', os.path.basename(cqu_ado_smi_file).lower())
        os.makedirs(os.path.dirname(dst_cqu_ado_smi_file), exist_ok=True)
        shutil.copyfile(cqu_ado_smi_file, dst_cqu_ado_smi_file)

    if os.path.isfile(dst_protein_file) or os.path.isfile(dst_smiles_file) \
            or os.path.isfile(dst_cqu_pocket_10A_file) or os.path.isfile(dst_cqu_ado_sdf_file) or os.path.isfile(dst_cqu_ado_smi_file):
        logger.info(f'任务ID：{task_id}, 蛋白质文件：{dst_protein_file}, 化合物文件：{dst_smiles_file}')
        return dst_protein_file, dst_smiles_file, dst_cqu_pocket_10A_file, dst_cqu_ado_sdf_file, dst_cqu_ado_smi_file
    else:
        logger.info(f'任务ID：{task_id}, 文件创建失败，没有任何文件被上传 ！！！！！！！！！！！！！!')
        gr.Error(" 文件创建失败，没有任何文件被上传  💥!", duration=5)
        return None, None, None, None, None


def get_protein_ligand_pocket(input_protein_file_path):
    pdb_id = os.path.basename(input_protein_file_path)[:-4]
    ligand_save_dir = os.path.splitext(input_protein_file_path)[0]
    print('protein格式cif转为pdb')
    output_protein_file_path = os.path.join(ligand_save_dir, f'{pdb_id}.pdb')
    ligand_format_conversion(input_protein_file_path, output_protein_file_path)
    print('提取蛋白质文件中所有的chain，并分类保存到指定文件夹')
    input_protein_file_path = os.path.join(ligand_save_dir, f'{pdb_id}.pdb')  # 指定PDB文件路径（请替换为你的实际文件路径）
    save_chain_dir = os.path.join(ligand_save_dir, 'chains')  # 输出目录名称
    get_chain_class(input_protein_file_path, save_chain_dir)
    # =====================================================================
    print('获取排序最高的蛋白质文件路径')
    protein_path_list = glob(os.path.join(save_chain_dir, '*_蛋白质_*.pdb'))
    if len(protein_path_list) == 0:
        gr.Error("未能提取蛋白质链 💥!", duration=60)
    ligand_path_list = glob(os.path.join(save_chain_dir, '*_配体_*.pdb'))
    if len(ligand_path_list) == 0:
        gr.Error("未能提取配体 💥!", duration=60)
    protein_pocket_size_list = []
    for protein_pdb in protein_path_list:
        for ligand_x in ligand_path_list:
            try:
                ligand_add_h_path = f"{os.path.join(save_chain_dir, os.path.basename(ligand_x)[:-4])}_ligand_addh.pdb"
                add_hydrogens(ligand_x, ligand_add_h_path)
                reference_ligand_mol2 = os.path.join(save_chain_dir, os.path.basename(ligand_x)[:-4]) + '.mol2'
                reference_ligand_sdf = os.path.join(save_chain_dir, os.path.basename(ligand_x)[:-4]) + '.sdf'
                ligand_format_conversion(ligand_add_h_path, reference_ligand_mol2)
                ligand_format_conversion(ligand_add_h_path, reference_ligand_sdf)
                pocket_pdb = os.path.join(save_chain_dir, os.path.basename(ligand_x)[:-4]) + '_pocket_ligH10A.pdb'
                get_pocket(reference_ligand_sdf, protein_pdb, reference_ligand_mol2, pocket_pdb, exwithin=10)
                protein_pocket_size_list.append([os.path.getsize(pocket_pdb), protein_pdb, ligand_x])
            except:
                pass
            for del_path in [ligand_add_h_path, reference_ligand_mol2, reference_ligand_sdf, pocket_pdb]:
                if os.path.isfile(del_path):
                    os.remove(del_path)
    protein_pocket_size_sorted = sorted(protein_pocket_size_list, key=lambda x: x[0])
    pocket_size, protein_max_path, ligand_max_path = protein_pocket_size_sorted[-1]
    # =================================================================
    print('删除旧氢，添加新氢')
    ligand_add_h_path = f"{os.path.join(ligand_save_dir, pdb_id)}_ligand_addh.pdb"
    add_hydrogens(ligand_max_path, ligand_add_h_path)
    print('复制文件到结构性目录中')
    reference_ligand_sdf = f"{os.path.join(ligand_save_dir, 'glide_pos', pdb_id)}_ligand.sdf"
    reference_ligand_mol2 = f"{os.path.join(ligand_save_dir, 'protein', pdb_id)}_ligand.mol2"
    protein_pdb = f"{os.path.join(ligand_save_dir, 'protein', pdb_id)}_protein.pdb"
    ligand_format_conversion(ligand_add_h_path, reference_ligand_sdf)
    ligand_format_conversion(ligand_add_h_path, reference_ligand_mol2)
    shutil.copy(protein_max_path, protein_pdb)
    print('提取配体口袋')
    pocket_pdb_12A = f"{os.path.join(ligand_save_dir, 'protein', pdb_id)}_pocket_ligH12A.pdb"
    get_pocket(reference_ligand_sdf, protein_pdb, reference_ligand_mol2, pocket_pdb_12A, exwithin=12)
    pocket_pdb_10A = f"{os.path.join(ligand_save_dir, 'protein', pdb_id)}_pocket_ligH10A.pdb"
    get_pocket(reference_ligand_sdf, protein_pdb, reference_ligand_mol2, pocket_pdb_10A, exwithin=10)
    # 配体化合物sdf、配体化合物mol2、蛋白质源文件不含配体pdb、口袋文件不含配体
    logger.info(f'预处理已完成，保存地址为：{ligand_save_dir}')
    print(reference_ligand_sdf, reference_ligand_mol2, protein_pdb, pocket_pdb_12A, pocket_pdb_10A)
    return reference_ligand_sdf, reference_ligand_mol2, protein_pdb, pocket_pdb_12A, pocket_pdb_10A


# gr.Error("An error occurred 💥!", duration=5)
# gr.Info("Helpful info message ℹ️", duration=5)
# gr.Warning("A warning occured ⛔️!", duration=5)
def run_main(model_name_list, protein_file, smiles_file, cqu_pocket_10A_file, cqu_ado_sdf_file, cqu_ado_smi_file):
    task_id = test_create_task_id()
    logger.info(f'开始对接任务，任务ID：{task_id}')
    dst_protein_file, dst_smiles_file, dst_cqu_pocket_10A_file, dst_cqu_ado_sdf_file, dst_cqu_ado_smi_file = \
        copy_input_file2out(task_id, protein_file, smiles_file, cqu_pocket_10A_file, cqu_ado_sdf_file, cqu_ado_smi_file)
    cif_path = None
    if dst_protein_file.endswith('.cif'):
        cif_path = dst_protein_file
    # 验证smi的合理性
    csv_df = pd.read_csv(dst_smiles_file)
    df_valus_np = csv_df.values
    verify_smiles_list = get_verify_smiles(df_valus_np)
    smi_df = pd.DataFrame(verify_smiles_list, columns=csv_df.columns.tolist())
    smi_df.to_csv(dst_smiles_file, index=False)  # index=False不补首列序列
    try:
        # 配体化合物sdf、配体化合物mol2、蛋白质源文件不含配体pdb、口袋文件不含配体pocket_ligH12A
        reference_ligand_sdf, reference_ligand_mol2, protein_pdb, pocket_pdb_12A, pocket_pdb_10A = get_protein_ligand_pocket(dst_protein_file)
    except Exception as e:
        logger.error(f'{task_id}, {e}')
        print(task_id, e)
        gr.Error(f"蛋白、配体分离报错：{e}", duration=60)
        return f'当前任务ID：{task_id},蛋白、配体分离报错：{e}'

    logger.info(f'预处理数据完成，任务ID：{task_id}，{reference_ligand_sdf}')
    logger.info(f'预处理数据完成，任务ID：{task_id}，{reference_ligand_mol2}')
    logger.info(f'预处理数据完成，任务ID：{task_id}，{protein_pdb}')
    logger.info(f'预处理数据完成，任务ID：{task_id}，{pocket_pdb_10A}')
    logger.info(f'预处理数据完成，任务ID：{task_id}，{pocket_pdb_12A}')
    rmq_args_dict = {
        'task_id': task_id, 'pocket_ligH10A_path': pocket_pdb_10A, 'pocket_ligH12A_path': pocket_pdb_12A, 'ligand_sdf_path': reference_ligand_sdf,
        'smiles_file': dst_smiles_file, 'ligand_mol2_path': reference_ligand_mol2, 'cif_path': cif_path,
        'dst_cqu_pocket_10A_file': dst_cqu_pocket_10A_file, 'dst_cqu_ado_sdf_file': dst_cqu_ado_sdf_file, 'dst_cqu_ado_smi_file': dst_cqu_ado_smi_file
    }
    rmq_args_dict_str = json.dumps(rmq_args_dict)
    for model_name in model_name_list:
        if model_name == "model_PG":  # PhoreGen
            rmq_sent(rmq_args_dict_str, queue_name='phore_gen')
        if model_name == "model_SG":  # ScaffoldGVAE
            rmq_sent(rmq_args_dict_str, queue_name='scaffold_gvae')
        if model_name == "model_MC":  # MolCRAFT
            rmq_sent(rmq_args_dict_str, queue_name='mol_craft')
        if model_name == "model_TG":  # TamGen
            rmq_sent(rmq_args_dict_str, queue_name='tam_gen')
        # ===================重大
        if model_name == "edu_cqu_ADO":  # EduCquAdo
            rmq_sent(rmq_args_dict_str, queue_name='edu_cqu_ado')
        # ===================重邮
    # 记录生成模型
    with open(os.path.join(cfg['generate_result_path'], task_id, 'info_file.json'), "w+", encoding='utf-8') as json_f:
        info_dict = {'model_list': model_name_list}
        json.dump(info_dict, json_f, ensure_ascii=False, indent=4)
    return f'当前任务ID：{task_id}'


def get_task_progress(task_id):
    task_parent_dir = os.path.join(cfg['generate_result_path'], task_id)
    with open(os.path.join(cfg['generate_result_path'], task_id, 'info_file.json')) as json_file:
        json_lines = json_file.readlines()
    json_data = json.loads(''.join(json_lines))
    model_name_list = json_data['model_list']
    log_info = ''
    smiles_file_list = []
    model_dict = cfg['model']
    for model_key in model_name_list:
        model_name = model_dict[model_key]
        try:
            if model_name == 'PhoreGen':
                result_dir = os.path.join(task_parent_dir, 'ResultPhoreGen')
                with open(os.path.join(result_dir, 'log.txt'), encoding='utf-8') as log_file:
                    log_lines = log_file.readlines()

                log_info += ''.join(log_lines) + '\r\n'
                smiles_file_list += glob(os.path.join(result_dir, '*_SMILES_all.txt'))

            if model_name == 'ScaffoldGVAE':
                result_dir = os.path.join(task_parent_dir, 'ResultScaffoldGVAE')
                with open(os.path.join(result_dir, 'log.txt'), encoding='utf-8') as log_file:
                    log_lines = log_file.readlines()
                log_info += ''.join(log_lines) + '\r\n'
                smiles_file_list += glob(os.path.join(result_dir, '*_sample_*.txt'))
            if model_name == "MolCRAFT":
                result_dir = os.path.join(task_parent_dir, 'ResultMolCRAFT')
                with open(os.path.join(result_dir, 'log.txt'), encoding='utf-8') as log_file:
                    log_lines = log_file.readlines()
                log_info += ''.join(log_lines) + '\r\n'
                smiles_file_list += glob(os.path.join(result_dir, 'model_MC_sample_result.zip'))
            if model_name == "TamGen":
                result_dir = os.path.join(task_parent_dir, 'ResultTamGen')
                with open(os.path.join(result_dir, 'log.txt'), encoding='utf-8') as log_file:
                    log_lines = log_file.readlines()
                log_info += ''.join(log_lines) + '\r\n'
                smiles_file_list += glob(os.path.join(result_dir, '*/*_flatten.tsv'))
            if model_name == "EduCquAdo":
                result_dir = os.path.join(task_parent_dir, 'ResultEduCquAdo')
                with open(os.path.join(result_dir, 'log.txt'), encoding='utf-8') as log_file:
                    log_lines = log_file.readlines()
                log_info += ''.join(log_lines) + '\r\n'
                smiles_file_list += glob(os.path.join(result_dir, 'model_cqu_ado_generate_result.zip'))
        except Exception as e:
            logger.error(f'{task_id}\t{e}')
    return log_info, smiles_file_list
